import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QGridLayout, QLineEdit, QPushButton
import psycopg2


class MyApp(QWidget):

    def __init__(self):
        super().__init__()

        # 初始化界面
        self.initUI()

        # 初始化数据库
        self.initDB()

    def initUI(self):
        # 设置窗口标题
        self.setWindowTitle('OpenGauss查询工具')
        # 设置窗口大小
        self.resize(300, 150)
        # 创建文本标签和输入框，用于输入数据库登录信息
        hostLabel = QLabel('主机名：')
        portLabel = QLabel('端口号：')
        dbLabel = QLabel('数据库：')
        userLabel = QLabel('用户名：')
        passwordLabel = QLabel('密   码：')
        self.hostEdit = QLineEdit()
        self.portEdit = QLineEdit()
        self.dbEdit = QLineEdit()
        self.userEdit = QLineEdit()
        self.passwordEdit = QLineEdit()
        self.passwordEdit.setEchoMode(QLineEdit.Password)

        # 创建查询条件输入框和查询按钮
        self.queryEdit = QLineEdit()
        queryBtn = QPushButton('查询', self)
        queryBtn.clicked.connect(self.onQueryClick)

        # 创建网格布局，并将文本标签和输入框添加到布局中
        grid = QGridLayout()
        grid.setSpacing(10)
        grid.addWidget(hostLabel, 1, 0)
        grid.addWidget(self.hostEdit, 1, 1)
        grid.addWidget(portLabel, 2, 0)
        grid.addWidget(self.portEdit, 2, 1)
        grid.addWidget(dbLabel, 3, 0)
        grid.addWidget(self.dbEdit, 3, 1)
        grid.addWidget(userLabel, 4, 0)
        grid.addWidget(self.userEdit, 4, 1)
        grid.addWidget(passwordLabel, 5, 0)
        grid.addWidget(self.passwordEdit, 5, 1)
        grid.addWidget(self.queryEdit, 6, 0, 1, 2)
        grid.addWidget(queryBtn, 7, 1)

        # 将网格布局应用到窗口中
        self.setLayout(grid)

    def initDB(self):
        # 连接OpenGauss数据库
        self.conn = psycopg2.connect(
            host=self.hostEdit.text(),
            port=self.portEdit.text(),
            database=self.dbEdit.text(),
            user=self.userEdit.text(),
            password=self.passwordEdit.text()
        )
        # 创建游标
        self.cur = self.conn.cursor()

        # 创建加密密钥
        self.cur.execute("SELECT pgp_sym_encrypt('mykey', 'mykey')")
        self.key = self.cur.fetchone()[0].decode('utf-8')

    def onQueryClick(self):
        # 获取查询条件
        condition = self.queryEdit.text()

        # 判断是否有查询条件
        if condition:
            # 对查询条件进行加密，并将查询语句发送到数据库进行执行
            encrypt_condition = self.cur.mogrify(
                "SELECT * FROM mytable WHERE Equal(salary, Decrypt(%s, 'mykey'))",
                (psycopg2.extensions.QuotedString(condition.encode('utf-8')).getquoted(),)
            )
            self.cur.execute(encrypt_condition)

            # 获取查询结果
            rows = self.cur.fetchall()
            # 显示解密后的查询结果
            if rows:
                for row in rows:
                    row = list(row)
                    row[1] = self.cur.mogrify(
                        "Decrypt(%s, 'mykey')",
                        (row[1],)
                    )
                    self.queryEdit.insertPlainText(
                        ','.join([str(cell, encoding='utf-8') for cell in row]) + '\n'
                    )
            else:
                self.queryEdit.clear()

    def closeEvent(self, event):
        # 关闭游标和连接
        self.cur.close()
        self.conn.close()


if __name__ == '__main__':
    # 创建并应用QApplication
    app = QApplication(sys.argv)
    ex = MyApp()
    ex.show()
    sys.exit(app.exec_())

